package org.zookeeper.examples.queue.fifo;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.TreeMap;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooDefs.Ids;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.ACL;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 基于zookeeper的分布式fifo队列 queue队列方法:
 *
 * @author slydm
 */
public class ZFIFOQueue {

  private static final Logger logger = LoggerFactory.getLogger(ZFIFOQueue.class);
  public static final String QUEUE_ZNODE_PREFIX = "/";
  public static final String QUEUE_ZNODE_SUFFIX = "zfifo-queue";
  public static final String ELEMENT_PREFIX = "el-";

  private ZooKeeper zookeeper;
  private String rootZnode;
  private String childSequelZnode;
  private ArrayList<ACL> acl = Ids.OPEN_ACL_UNSAFE;

  public ZFIFOQueue(ZooKeeper zookeeper, String prefix) {
    this.zookeeper = zookeeper;
    this.rootZnode = (null == prefix ? QUEUE_ZNODE_PREFIX + QUEUE_ZNODE_SUFFIX
        : prefix.endsWith("/") ? prefix + QUEUE_ZNODE_SUFFIX : prefix + "/" + QUEUE_ZNODE_SUFFIX);
    this.childSequelZnode = rootZnode + "/" + ELEMENT_PREFIX;
  }

  /**
   * 将元素添加队列尾部 {@link java.util.concurrent.locks.AbstractQueuedSynchronizer#}中的enq()方法
   *
   * @param data
   * @return
   * @throws KeeperException
   * @throws InterruptedException
   */
  public boolean offer(byte[] data) throws KeeperException, InterruptedException {

    for (; ; ) {
      try {
        zookeeper.create(childSequelZnode, data, acl, CreateMode.PERSISTENT_SEQUENTIAL);
        return true;
      } catch (KeeperException.NoNodeException e) {
        zookeeper.create(rootZnode, null, acl, CreateMode.PERSISTENT);
      }
    }
  }


  /**
   * 返回队列第一个元素
   *
   * @return 队列第一个元素, 队列空则null
   * @throws KeeperException
   * @throws InterruptedException
   */
  public byte[] peek() throws KeeperException, InterruptedException {

    Map<Long, String> orderChildren = orderChildren();

    Iterator<Entry<Long, String>> iterator = orderChildren.entrySet().iterator();

    while (iterator.hasNext()) {
      String znode = iterator.next().getValue();
      try {
        return zookeeper.getData(znode, false, null);
      } catch (KeeperException.NoNodeException e) {
        // 该节点被其它客户端删除, 不用处理, 继续找下一个
      }
    }
    return null;
  }

  /**
   * 删除队列头部元素并返回
   *
   * @return 队列第一个元素, 队列空则null
   */
  public byte[] poll() throws KeeperException, InterruptedException {

    Map<Long, String> orderChildren = orderChildren();

    Iterator<Entry<Long, String>> iterator = orderChildren.entrySet().iterator();

    while (iterator.hasNext()) {
      String znode = iterator.next().getValue();
      try {
        byte[] data = zookeeper.getData(znode, false, null);
        zookeeper.delete(znode, -1);
        return data;
      } catch (KeeperException.NoNodeException e) {
        // 该节点被其它客户端删除, 不用处理, 继续找下一个
      }
    }
    return null;
  }


  private Map<Long, String> orderChildren() throws KeeperException, InterruptedException {

    List<String> childrenNames = zookeeper.getChildren(rootZnode, null);

    Map<Long, String> orderChildrens = new TreeMap<>();
    for (String childrenName : childrenNames) {
      boolean isValidChildrenName = childrenName
          .regionMatches(0, ELEMENT_PREFIX, 0, ELEMENT_PREFIX.length());
      if (!isValidChildrenName) {
        logger.warn("找到不符合命名规范的节点: {}", childrenName);
        continue;
      }

      try {
        Long id = Long.parseLong(childrenName.substring(ELEMENT_PREFIX.length()));
        orderChildrens.put(id, rootZnode + "/" + childrenName);
      } catch (NumberFormatException e) {
        logger.warn("找到不符合命名规范的节点: {}", childrenName, e);
      }
    }

    return orderChildrens;
  }


  public String getRootZnode() {
    return rootZnode;
  }

  public String getChildSequelZnode() {
    return childSequelZnode;
  }
}
